from genshi.core import Attrs, QName, stripentities
from genshi.core import END, START, TEXT, COMMENT
from genshi import XML



class Replacer(object):
    def __init__(self,tag,stream):
        self.tag = tag
        self.stream = list(stream)
    def __call__(self,stream):
        for kind, data, pos in stream:
            if kind is START:
                tag, attrs = data
                if tag == self.tag:
                    #temp = toc.generate()
                    for kind, data, pos in self.stream:
                        yield kind, data, pos
                    continue
            elif kind is END:
                tag = data
                if tag == self.tag:
                    continue
            yield kind, data, pos            

class TOCMaker(object):
    def __init__(self,depth):
        self.depth = depth
        self.levels = ['h'+str(l) for l in range(1,depth+1)]
    def __call__(self,stream):
        yield START, (QName('ol'),Attrs([(QName('class'),'toc'),])),(None,-1,-1)
        current_level = 1
        in_heading = False
        first_item_done = False
        for kind,data,pos in stream:
            if kind is START:
                tag, attrs = data
                if tag in self.levels:
                    in_heading = True
                    #if not current_level:
                    #    current_level = tag
                    level = int(tag[1])
                    if level > current_level:
                        if not first_item_done:
                            yield START, (QName('li'),Attrs()),pos
                        yield START, (QName('ol'),Attrs()),pos
                        for x in range(level-current_level-1):
                            yield START, (QName('li'),Attrs()),pos
                            yield START, (QName('ol'),Attrs()),pos
                    elif level < current_level:
                        yield END, QName('li'),pos
                        for x in range(current_level-level):
                            yield END, QName('ol'),pos
                            yield END, QName('li'),pos
                    elif level == current_level and first_item_done:
                        yield END, QName('li'),pos
                    first_item_done = True
                    current_level = int(tag[1])
                    yield START, (QName('li'),Attrs()),pos
                    id_ = attrs[0][1]
                    href = '#'+id_
                    yield START, (QName('a'),Attrs([(QName('href'),href)])),pos
                    continue
                elif tag == 'a' and in_heading:
                    continue
            elif kind is END:
                tag = data
                if tag in self.levels:
                    in_heading = False
                    yield END, QName('a'),pos
                    yield TEXT, '\n', pos
                    #yield END, QName('li'),pos
                    continue
                elif tag == 'a' and in_heading:
                    continue
            if in_heading:
                yield kind,data,pos
        yield END, QName('li'),pos        
        if current_level > 1:
            for x in range(current_level-1):
                yield END, QName('ol'),pos
                yield END, QName('li'),pos
        yield END, QName('ol'),pos
           

if __name__ == "__main__":

    stream = XML("""
<div>
<toc />
<h2 id="!level-2"><a href="bogus">level 2.1</a> 2</h2>
<h2 id="!level-2_1">level 2.2</h2> 
<h1 id="!level-1-1">level 1.1</h1>
<h2 id="!level-2_2">level 2.3</h2>
<h1 id="!level-1-2">level 1.2</h1></div>""")
    stream = stream.filter(Replacer('toc', stream.filter(TOCMaker(3))))
    #stream = XML(stream.render())
    output = stream.render()
    assert output == """\
<div>
<ol class="toc"><li><ol><li><a href="#!level-2">level 2.1 2</a>
</li><li><a href="#!level-2_1">level 2.2</a>
</li></ol></li><li><a href="#!level-1-1">level 1.1</a>
<ol><li><a href="#!level-2_2">level 2.3</a>
</li></ol></li><li><a href="#!level-1-2">level 1.2</a>
</li></ol>
<h2 id="!level-2"><a href="bogus">level 2.1</a> 2</h2>
<h2 id="!level-2_1">level 2.2</h2>
<h1 id="!level-1-1">level 1.1</h1>
<h2 id="!level-2_2">level 2.3</h2>
<h1 id="!level-1-2">level 1.2</h1></div>"""

    stream = XML("""
<div>
<toc />
<h1 id="!level-1-1">level 1.1</h1>
<h2 id="!level-2_2">level 2.3</h2>
<h3 id="!level-1-2">level 1.2</h3>
</div>""")
    stream = stream.filter(Replacer('toc', stream.filter(TOCMaker(3))))
    #print stream
    #stream = XML(stream.render())
    output = stream.render()

    assert output == """\
<div>
<ol class="toc"><li><a href="#!level-1-1">level 1.1</a>
<ol><li><a href="#!level-2_2">level 2.3</a>
<ol><li><a href="#!level-1-2">level 1.2</a>
</li></ol></li></ol></li></ol>
<h1 id="!level-1-1">level 1.1</h1>
<h2 id="!level-2_2">level 2.3</h2>
<h3 id="!level-1-2">level 1.2</h3>
</div>"""

    

    
    
